#!/usr/bin/python3
# Prof. Shi Lei's Dynamic Programming Algorithm
# programmed by Sun Smallwhite <niasw@pku.edu.cn>(https://github.com/niasw)

import numpy
import matplotlib
matplotlib.use('Agg'); # plot and save, never showing

import matplotlib.pyplot as plt

def calcXvector(clussprs):
  '''
# for a paper p:
#  x(p)=1/sqrt(|#c(p)|)
# input:
#  clussprs: cluster seperators: [0,p_1,p_2,...,p_n,|#P|]
  '''
  xvec=numpy.zeros(clussprs[-1])+1.;
  for c in range(1,len(clussprs)):
   xvec[range(clussprs[c-1],clussprs[c])]=1./numpy.sqrt(clussprs[c]-clussprs[c-1]);
  return xvec;

def monitorXandErr(eigvec,clussprs):
  '''
# monitoring x vector and error function
  '''
  xvec=calcXvector(clussprs);
  reerr=1-xvec.dot(eigvec)**2/xvec.dot(xvec)/eigvec.dot(eigvec);
  return xvec,reerr;

def visualizeDivide(eigvec,clussprs):
  '''
# visualize dividing progress
  '''
  mappedxvec=calcXvector(clussprs);
  mappedeigv=eigvec;
  maponexvec=mappedxvec/numpy.sqrt(len(clussprs)-1);
  plt.clf();
  plt.plot(range(0,len(eigvec)),mappedeigv,color="green",linewidth=2,linestyle='--',label='eigenvector');
  plt.plot(range(0,len(eigvec)),maponexvec,color="blue", linewidth=2,linestyle='-',label='x vector');
  plt.grid();
  plt.xlabel('paper index (sorted by eigvec)');
  plt.ylabel('vector component');
  plt.legend(loc='upper left');
  plt.savefig('../out/xveig'+str(len(clussprs)-1)+'.png');
  # plt.show(); # don't forget to commentize "matplotlib.use('Agg');" in the beginning

def clussprs2clusters(clussprs,sortmap):
  '''
# generate clusters set from cluster seperator data
  '''
  def _idxfunc(x):
    return x[1];
  clusters=[list(set([_idxfunc(sortmap[p]) for p in range(clussprs[c],clussprs[c+1])])) for c in range(0,len(clussprs)-1)];
  return clusters;

def calcClusters(eigvec, clusnum=20, monitor=False, outputall=False):
  '''
# set paper clusters to make b=sum(a|c)/sqrt(#c) approching sum(a)/sum(sqrt(#c))
# Algorithm by Shi Lei <shijim@gmail.com>
# Program by Sun Smallwhite <niasw@pku.edu.cn>(https://github.com/niasw)

# input:
#  eigenvector
#  clusnum: max number of clusters
#  monitor: whether visualize the progress
#  outputall: output results when clusnum is lower than setting value
# output:
#  outputall=False: snapshot (last history state detail): clusters set [[paper indices],...]
#  outputall=True: history: [cluster seperators when clusnum=2, cluster seperators when clusnum=3, ..., cluster seperators when clusnum=clusnum], sortmap: sort eigvec from low to high
  '''
  dimension=len(eigvec);
  history=[]; # record all best division with lower clusnum

  def _keyfunc(x):
   return x[0];
  sortmap=sorted([(eigvec[p],p) for p in range(0,len(eigvec))],key=_keyfunc); # (eigen_value,origin_order)
  sortbck=sorted([(sortmap[p][1],p) for p in range(0,len(eigvec))],key=_keyfunc); # (origin_order,eigen_order)
  mappedeigv=numpy.float64([e[0] for e in sortmap]);
  numpyeigvec=numpy.array(mappedeigv); # different from list when slicing

  if (monitor):
   print('Dynamic Programming Algorithm:'+'\n  paper number: '+str(dimension)+'\n  set cluster number: '+str(clusnum));
  if (dimension<clusnum):
   print('Setting Cluster Number is smaller than Current Cluster Number, reset to equal.');
   clusnum=dimension;

  # Best 2D array initialization: http://stackoverflow.com/a/2398187
  clussprs=[x[:] for x in [[None]*dimension]*clusnum]; # best seperators (considering recursive relation, only last seperator recorded)
  clusobjs=[x[:] for x in [[0]*dimension]*clusnum]; # best objectives
  eigvsums=[None]*dimension; # eigenvector component sum

  # boundary conditions
  eigvtmp=0;
  for kminus1 in range(0,clusnum):
   eigvtmp+=numpyeigvec[kminus1];
   eigvsums[kminus1]=eigvtmp;
   clussprs[kminus1][kminus1]=kminus1;
   clusobjs[kminus1][kminus1]=eigvtmp;
   clussprs[0][kminus1]=0;
   clusobjs[0][kminus1]=eigvtmp/numpy.sqrt(kminus1+1);
  for jminus1 in range(clusnum,dimension):
   eigvtmp+=numpyeigvec[jminus1];
   eigvsums[jminus1]=eigvtmp;
   clussprs[0][jminus1]=0;
   clusobjs[0][jminus1]=eigvtmp/numpy.sqrt(jminus1+1);

  for kminus1 in range(1,clusnum):
   if monitor:
    print('Calculating situation: clusnum='+str(kminus1+1));
   for jminus1 in range(kminus1+1,dimension):
    tmp=[(clusobjs[kminus1-1][jminus1-m]+(eigvsums[jminus1]-eigvsums[jminus1-m])/numpy.sqrt(m)) for m in range(1,jminus1-kminus1+2)];
    tmpidx=numpy.argmax(tmp);
    clussprs[kminus1][jminus1]=jminus1-tmpidx;
    clusobjs[kminus1][jminus1]=tmp[tmpidx];

   if (monitor or outputall):
    result=[dimension];
    for it in range(0,kminus1+1):
      tmpidx=clussprs[kminus1-it][result[0]-1];
      result.insert(0,tmpidx);
    if monitor:
      xvec,reerr=monitorXandErr(mappedeigv,list(result));
      visualizeDivide(mappedeigv,list(result));
      print('unparallel ratio = '+str(reerr));
    if outputall:
      history.append(clussprs2clusters(list(result),sortmap));
  result=[dimension];
  for kminus1 in range(0,clusnum):
    tmpidx=clussprs[clusnum-1-kminus1][result[0]-1];
    result.insert(0,tmpidx);
  if outputall:
    history.append(clussprs2clusters(list(result),sortmap));
    return history;
  else:
    clusters=clussprs2clusters(list(result),sortmap);
    return clusters;

